from langchain.prompts import PromptTemplate, FewShotPromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector
from langchain_core.output_parsers import StrOutputParser

from create_llm import create_llm

# 示例库（包含不同长度的示例）
examples = [
  {"question": "天气如何？", "answer": "我无法实时获取天气，建议查看天气预报网站。"},
  {"question": "北京有哪些著名景点？", "answer": "北京的著名景点包括故宫、长城、颐和园等。"},
  {"question": "详细介绍Python的面向对象编程", "answer": "Python的面向对象编程支持类、继承、多态...（此处省略200字）"}
]

example_prompt = PromptTemplate(
  input_variables=["question", "answer"],
  template="问题: {question}\n回答: {answer}"
)

# 创建长度选择器（最大长度=100 tokens）
example_selector = LengthBasedExampleSelector(
  examples=examples,
  example_prompt=example_prompt,
  max_length=30  # 根据模型上下文窗口调整
)

# 创建少样本提示模板
prompt = FewShotPromptTemplate(
  example_selector=example_selector,
  example_prompt=example_prompt,
  prefix="请回答用户问题:",
  suffix="问题: {input}\n回答:",
  input_variables=["input"]
)

print("短问题输入，提示词会包含更多示例============>>>>>>>>>>>>>")
# 短输入：会包含更多示例
print(prompt.format(input="上海有哪些景点？"))

# 短问题输入，提示词会包含更多示例============>>>>>>>>>>>>>
# 请回答用户问题:
#
# 问题: 天气如何？
# 回答: 我无法实时获取天气，建议查看天气预报网站。
#
# 问题: 北京有哪些著名景点？
# 回答: 北京的著名景点包括故宫、长城、颐和园等。
#
# 问题: 详细介绍Python的面向对象编程
# 回答: Python的面向对象编程支持类、继承、多态...（此处省略200字）
#
# 问题: 上海有哪些景点？
# 回答:

print("\n长问题输入，提示词会包含更少示例============>>>>>>>>>>>>>")
# 长输入：会自动减少示例
print(prompt.format(input="""
请详细介绍 Python 的异步编程模型，包括 asyncio 库的使用方法、协程原理以及与多线程的对比。Python 的异步编程模型是一种高效处理 IO 密集型任务的编程范式，核心依托于协程（coroutine）和事件循环（event loop），能够在单线程内实现并发操作，避免因等待 IO 操作而浪费 CPU 资源。
asyncio 库是 Python 异步编程的标准库，其使用方法主要围绕事件循环和协程展开。通过async def定义协程函数，在函数内部使用await关键字来暂停协程执行，让出控制权给事件循环，待 IO 操作完成后再恢复执行。例如，在进行网络请求时，await asyncio.sleep(1)模拟网络延迟，在此期间事件循环可调度其他协程执行，提升程序整体效率。asyncio.gather()方法则能批量执行多个协程任务，并等待所有任务完成后返回结果。
"""))

# 请回答用户问题:
#
# 问题: 天气如何？
# 回答: 我无法实时获取天气，建议查看天气预报网站。
#
# 问题:
# 请详细介绍 Python 的异步编程模型，包括 asyncio 库的使用方法、协程原理以及与多线程的对比。Python 的异步编程模型是一种高效处理 IO 密集型任务的编程范式，核心依托于协程（coroutine）和事件循环（event loop），能够在单线程内实现并发操作，避免因等待 IO 操作而浪费 CPU 资源。
# asyncio 库是 Python 异步编程的标准库，其使用方法主要围绕事件循环和协程展开。通过async def定义协程函数，在函数内部使用await关键字来暂停协程执行，让出控制权给事件循环，待 IO 操作完成后再恢复执行。例如，在进行网络请求时，await asyncio.sleep(1)模拟网络延迟，在此期间事件循环可调度其他协程执行，提升程序整体效率。asyncio.gather()方法则能批量执行多个协程任务，并等待所有任务完成后返回结果。

# /*---------------------------------------用这个Prompt测试调用大模型-------------------------------------------*/

_llm = create_llm()

chain = prompt | _llm | StrOutputParser()

for chunk in chain.stream({"input": "北京有哪些景点"}):
  print(chunk, end="", flush=True)
